import os
from modules.utils import *
from modules.downloader import *
from modules.show import *
from modules.csv_downloader import *

from modules.utils import bcolors as bc
from glob import glob
import json


def save_args_conf(args, dataset_path):
	print(bc.INFO+ ' saving dataset configurations at '+ os.path.join(dataset_path, 'config.json')+ bc.ENDC)
	with open(os.path.join(dataset_path, 'config.json'), 'w') as f:
		f.truncate()
		json.dump(vars(args), f)



def bounding_boxes_images(args, DEFAULT_OID_DIR):


	if not args.Dataset:
		dataset_dir = os.path.join(DEFAULT_OID_DIR, 'Dataset')
		csv_dir = os.path.join(DEFAULT_OID_DIR, 'csv_folder')
	else:
		dataset_dir = os.path.join(DEFAULT_OID_DIR, args.Dataset)
		csv_dir = os.path.join(DEFAULT_OID_DIR, 'csv_folder')

	name_file_class = 'class-descriptions-boxable.csv'
	CLASSES_CSV = os.path.join(csv_dir, name_file_class)



	if args.command == 'downloader':

		save_args_conf(args, dataset_dir)


		logo(args.command)

		if args.type_csv is None:
			print(bc.FAIL + 'Missing type_csv argument.' + bc.ENDC)
			exit(1)
		if args.classes is None:
			print(bc.FAIL + 'Missing classes argument.' + bc.ENDC)
			exit(1)
		if args.multiclasses is None:
			args.multiclasses = 0

		folder = ['train', 'validation', 'test']
		file_list = ['train-annotations-bbox.csv', 'validation-annotations-bbox.csv', 'test-annotations-bbox.csv']

		if args.classes[0].endswith('.txt'):
			with open(args.classes[0]) as f:
				args.classes = f.readlines()
				args.classes = [x.strip() for x in args.classes]
		else:
				args.classes = [arg.replace('_', ' ') for arg in args.classes]

		if args.multiclasses == '0':

			mkdirs(dataset_dir, csv_dir, args.classes, args.type_csv)

			for classes in args.classes:

				print(bc.INFO + 'Downloading {}.'.format(classes) + bc.ENDC)
				class_name = classes

				error_csv(name_file_class, csv_dir, args.yes)
				df_classes = pd.read_csv(CLASSES_CSV, header=None)

				try:
					class_code = df_classes.loc[df_classes[1].str.lower() == class_name.lower()].values[0][0]
				except:
					print(bc.FAIL+ " '"+class_name+"', please check spelling ! or that the class is already exist in the dataset"+ bc.ENDC)
					exit(1)


				if args.type_csv == 'train':
					name_file = file_list[0]
					df_val = TTV(csv_dir, name_file, args.yes)
					if not args.n_threads:
						download(args, df_val, folder[0], dataset_dir, class_name, class_code, class_list_for_yolo=args.classes)
					else:
						download(args, df_val, folder[0], dataset_dir, class_name, class_code, threads = int(args.n_threads), class_list_for_yolo=args.classes)

				elif args.type_csv == 'validation':
					name_file = file_list[1]
					df_val = TTV(csv_dir, name_file, args.yes)
					if not args.n_threads:
						download(args, df_val, folder[1], dataset_dir, class_name, class_code,class_list_for_yolo=args.classes)
					else:
						download(args, df_val, folder[1], dataset_dir, class_name, class_code, threads = int(args.n_threads), class_list_for_yolo=args.classes)

				elif args.type_csv == 'test':
					name_file = file_list[2]
					df_val = TTV(csv_dir, name_file, args.yes)
					if not args.n_threads:
						download(args, df_val, folder[2], dataset_dir, class_name, class_code, class_list_for_yolo=args.classes)
					else:
						download(args, df_val, folder[2], dataset_dir, class_name, class_code, threads = int(args.n_threads),class_list_for_yolo=args.classes)

				elif args.type_csv == 'all':
					for i in range(3):
						name_file = file_list[i]
						df_val = TTV(csv_dir, name_file, args.yes)
						if not args.n_threads:
 							download(args, df_val, folder[i], dataset_dir, class_name, class_code,class_list_for_yolo=args.classes)
						else:
							download(args, df_val, folder[i], dataset_dir, class_name, class_code, threads = int(args.n_threads),class_list_for_yolo=args.classes)
				else:
					print(bc.FAIL + 'csv file not specified' + bc.ENDC)
					exit(1)

		elif args.multiclasses == '1':

			class_list = args.classes
			print(bc.INFO + 'Downloading {} together.'.format(class_list) + bc.ENDC)
			multiclass_name = ['_'.join(class_list)]
			mkdirs(dataset_dir, csv_dir, multiclass_name, args.type_csv)

			error_csv(name_file_class, csv_dir, args.yes)
			df_classes = pd.read_csv(CLASSES_CSV, header=None)

			class_dict = {}
			for class_name in class_list:

				try:
					class_dict[class_name] = df_classes.loc[df_classes[1].str.lower() == class_name.lower()].values[0][0]
				except:
					print(bc.FAIL+ " '"+class_name+"', please check spelling ! or that the class is already exist in the dataset" + bc.ENDC)
					exit(1)

			for class_name in class_list:

				if args.type_csv == 'train':
					name_file = file_list[0]
					df_val = TTV(csv_dir, name_file, args.yes)
					if not args.n_threads:
						download(args, df_val, folder[0], dataset_dir, class_name, class_dict[class_name], class_list, class_list_for_yolo=args.classes)
					else:
						download(args, df_val, folder[0], dataset_dir, class_name, class_dict[class_name], class_list, int(args.n_threads), class_list_for_yolo=args.classes)

				elif args.type_csv == 'validation':
					name_file = file_list[1]
					df_val = TTV(csv_dir, name_file, args.yes)
					if not args.n_threads:
						download(args, df_val, folder[1], dataset_dir, class_name, class_dict[class_name], class_list, class_list_for_yolo=args.classes)
					else:
						download(args, df_val, folder[1], dataset_dir, class_name, class_dict[class_name], class_list, int(args.n_threads), class_list_for_yolo=args.classes)

				elif args.type_csv == 'test':
					name_file = file_list[2]
					df_val = TTV(csv_dir, name_file, args.yes)
					if not args.n_threads:
						download(args, df_val, folder[2], dataset_dir, class_name, class_dict[class_name], class_list, class_list_for_yolo=args.classes)
					else:
						download(args, df_val, folder[2], dataset_dir, class_name, class_dict[class_name], class_list, int(args.n_threads), class_list_for_yolo=args.classes)

				elif args.type_csv == 'all':
					for i in range(3):
						name_file = file_list[i]
						df_val = TTV(csv_dir, name_file, args.yes)
						if not args.n_threads:
							download(args, df_val, folder[i], dataset_dir, class_name, class_dict[class_name], class_list, class_list_for_yolo=args.classes)
						else:
							download(args, df_val, folder[i], dataset_dir, class_name, class_dict[class_name], class_list, int(args.n_threads), class_list_for_yolo=args.classes)


	elif args.command == 'visualizer':
		try:
			print(bc.INFO+ ' try to laod dataset configurations at '+ os.path.join(dataset_dir, 'config.json')+ bc.ENDC)
			with open(os.path.join(dataset_dir, 'config.json'), 'r') as f:
				config_dict = json.load(f)

			args.classes = config_dict['classes']
			args.multiclasses = config_dict['multiclasses']
			args.yoloLabelStyle = config_dict['multiclasses']

		except Exception as e:
			print(e)
			print(bc.FAIL+ ' cannot load Dataset configurations at '+ os.path.join(dataset_dir, 'config.json')+ bc.ENDC)
			print(bc.INFO+' you have to specifiy the dataset config by ur self if you faced any errors later' + bc.ENDC)

		logo(args.command)

		flag = 0

		while (True):
			if flag == 0:
				print("Which folder do you want to visualize (train, test, validation)? <exit>")
				image_dir = input("> ")
				flag = 1

				if image_dir == 'exit':
					exit(1)

				class_image_dir = os.path.join(dataset_dir, image_dir)

				print("Which class? <exit>")
				show_classes(os.listdir(class_image_dir))

				class_name = input("> ")
				if class_name == 'exit':
					exit(1)

			if args.multiclasses=='0':
				download_dir = os.path.join(dataset_dir, image_dir, class_name,'images')
				label_dir = os.path.join(dataset_dir, image_dir, class_name, 'labels')
			else:
				download_dir = os.path.join(dataset_dir, image_dir, class_name)
				label_dir = os.path.join(dataset_dir, image_dir, class_name)


			if not os.path.isdir(download_dir):
				print("[ERROR] Images folder not found")
				print("[INFO] If u downloaded your images in multiclasses format, then you have to specify it here also ")
				exit(1)
			if not os.path.isdir(label_dir):
				print("[ERROR] Labels folder not found")
				print("[INFO] If u downloaded your images in multiclasses format, then you have to specify it here also ")

				exit(1)

			index = 0

			images_files = glob(os.path.join(download_dir,'*.jpg'))
			images_files += glob(os.path.join(download_dir,'*.JPG'))
			images_files += glob(os.path.join(download_dir,'*.jpeg'))
			images_files += glob(os.path.join(download_dir,'*.JPEG'))
			images_files += glob(os.path.join(download_dir,'*.png'))
			images_files += glob(os.path.join(download_dir,'*.PNG'))



			print(dedent("""
                --------------------------------------------------------
                INFO:
                        - Press 'd' to select next image
                        - Press 'a' to select previous image
                        - Press 'e' to select a new class
                        - Press 'w' to select a new folder
                        - Press 'q' to exit
                  You can resize the window if it's not optimal
                --------------------------------------------------------
                """))
			classes_list = None
			if args.yoloLabelStyle:
				classes_list = args.classes

			index = show(class_name, download_dir, label_dir,images_files, index, args)

			while True:

				progression_bar(len(images_files), index+1)

				k = cv2.waitKey(0) & 0xFF

				if k == ord('d'):
					index += 1
				elif k == ord('a'):
					index -= 1
				
				elif k == ord('e'):
					cv2.destroyAllWindows()
					break
				elif k == ord('w'):
					flag = 0
					cv2.destroyAllWindows()
					break
				elif k == ord('q'):
					cv2.destroyAllWindows()
					exit(1)
					break
			
				cv2.destroyAllWindows()
				index =  show(class_name, download_dir, label_dir,images_files, index, args)


